import grpc
from concurrent import futures
import asr_service_pb2
import asr_service_pb2_grpc
import asyncio
import vosk  # 使用 Vosk 作为语音识别引擎

# 初始化 Vosk 模型
vosk.SetLogLevel(-1)
model = vosk.Model("vosk-model-small-en-us-daanzu-20200905")

class ASRService(asr_service_pb2_grpc.ASRServiceServicer):
    async def RecognizeStream(self, request_iterator, context):
        recognizer = vosk.KaldiRecognizer(model, 16000)
        for request in request_iterator:
            session_id = request.session_id
            audio_data = request.audio_data
            if recognizer.AcceptWaveform(audio_data):
                result = recognizer.Result()
                yield asr_service_pb2.RecognitionResult(text=result, session_id=session_id)
            else:
                partial_result = recognizer.PartialResult()
                yield asr_service_pb2.RecognitionResult(text=partial_result, session_id=session_id)

async def serve():
    server = grpc.aio.server(futures.ThreadPoolExecutor(max_workers=10))
    asr_service_pb2_grpc.add_ASRServiceServicer_to_server(ASRService(), server)
    listen_addr = "[::]:50051"
    server.add_insecure_port(listen_addr)
    print(f"Starting server on {listen_addr}")
    await server.start()
    await server.wait_for_termination()

if __name__ == "__main__":
    asyncio.run(serve())